// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
//

#include <kllm/openai/rerank.h>
#include <turbo/log/logging.h>

namespace kllm {

    turbo::Status format_aoi_slots_response(const KaiResponse &response, nlohmann::ordered_json &obj) {
        auto type = response.query_type();
        if(type == QUERY_SLOTS_SAVE) {

            obj = nlohmann::ordered_json{
                     {"id_slot",   response.slots_task().id_slot()},
                     {"filename",  response.slots_task().filename()},
                     {"n_saved",   response.slots_task().n_saved()}, // tokens saved
                     {"n_written", response.slots_task().n_written()},      // bytes written
                     {"timings",   {
                                           {"save_ms", response.slots_task().t_save_ms()}
                                   }}
             };
            return turbo::OkStatus();
        }
        if(type == QUERY_SLOTS_RESTORE) {
               obj = nlohmann::ordered_json{
                       {"id_slot",    response.slots_task().id_slot()},
                       {"filename",   response.slots_task().filename()},
                       {"n_restored", response.slots_task().n_restored()}, // tokens restored
                       {"n_read",     response.slots_task().n_read()},       // bytes read
                       {"timings",    {
                                              {"restore_ms", response.slots_task().t_restore_ms()}
                                      }}
               };
            return turbo::OkStatus();
        }

        if(type == QUERY_SLOTS_ERASE) {
            obj = nlohmann::ordered_json{
                    {"id_slot",  response.slots_task().id_slot()},
                    {"n_erased", response.slots_task().n_erased()}
            };
            return turbo::OkStatus();
        }
        return turbo::invalid_argument_error("bad query type for slots reault");
    }
}  // namespace kllm
